# -*- coding: utf-8 -*-
from DataSet import RGBImageSet
from model import ImageModel900
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
import numpy as np
import torch.nn as nn
import torch
import logging
import progressbar


logging.basicConfig(
    format='%(asctime)s: %(message)s',
    level=logging.INFO, filename=r'./log/logging_CNN_900_2.log'
)

# 显示的部件
widgets = [
    'Progress: ',
    progressbar.Percentage(), ' ',
    progressbar.Bar('#'), ' ',
    progressbar.Timer(), ' ',
    progressbar.ETA()
]


def showTransformedImg(img):
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.figure()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def test(net, imageset):
    # 训练集测试
    net.eval()
    imageset.trainMode()
    dataloader = DataLoader(imageset, batch_size=5,
                            num_workers=2, shuffle=False)
    eq, total = 0, 0
    for data in dataloader:
        images, labels = data
        images, labels = images.cuda(), labels.cuda()
        v = net(images)
        pred = torch.argmax(v, dim=1)
        compare = (labels == pred)
        eq += sum(compare).item()
        total += len(compare)
        # print(compare, eq, total)
    train_prec = eq / total

    imageset.testMode()
    dataloader = DataLoader(imageset, batch_size=5,
                            num_workers=2, shuffle=False)
    eq, total = 0, 0
    for data in dataloader:
        images, labels = data
        images, labels = images.cuda(), labels.cuda()
        v = net(images)
        pred = torch.argmax(v, dim=1)
        compare = (labels == pred)
        eq += sum(compare).item()
        total += len(compare)
    test_prec = eq / total
    return train_prec, test_prec


def train(epoch, lr=1e-3):
    img_dir = r'./dataset/Scaled'
    img_set = RGBImageSet(img_dir, inmemory=True, test_size=0.2)
    img_set.trainMode()
    dataloader = DataLoader(img_set, batch_size=5,
                            num_workers=2, shuffle=True)
    # pbar
    total = epoch * len(dataloader)
    pbar = progressbar.ProgressBar(widgets=widgets, maxval=total)
    pbar.start()
    # 训练
    net = ImageModel900()
    net = net.cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    step = 0
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
    # 
    last_loss = 1e10
    for ep in range(1, epoch + 1):
        logging.info('Epoch {}'.format(ep).center(30, '='))
        total_loss = 0
        # 进入训练模式
        net.train()
        img_set.trainMode()
        for i, data in enumerate(dataloader, 1):
            # 获取batch数据
            images, labels = data
            images, labels = images.cuda(), labels.cuda()
            # 梯度清零
            net.zero_grad()
            # forward
            outputs = net(images)
            # loss
            loss = criterion(outputs, labels)
            # backward
            loss.backward()
            # update
            # optimizer.step()
            scheduler.step()
            # print statistics
            loss_val = loss.item()
            total_loss += loss_val
            logging.info('Iter {}, loss = {:.3f}'.format(i, loss_val))
            pbar.update(step + 1)
            step += 1
        logging.info('Epoch {}, Total Loss: {:.3f}'.format(ep, total_loss))
        if ep % 10 == 0:
            train_prec, test_prec = test(net, img_set)
            logging.info('训练集: {:.3f}, 测试集: {:.3f}'.format(train_prec, test_prec))
            # 如果有必要就保存模型
            if total_loss < last_loss:
                torch.save(net, r'./save/scaled900_cnn_2.model')
                logging.info('Save Model At scaled900_cnn_2.model')
                last_loss = total_loss
    pbar.finish()


if __name__ == '__main__':
    train(epoch=500, lr=1e-2)
